nanopyx.core.io.zip_image_loader

 1import os
 2import zipfile
 3
 4import numpy as np
 5import skimage.io
 6from skimage import transform
 7
 8
 9class ZipTiffIterator:
10    _shape = None
11    _dtype = None
12    _im0 = None
13
14    def __init__(self, zip_file_path: str):
15        """
16        Iterator to sequentially open tiff files contained in a zip file
17        :param zip_file_path: path to the zip file
18        :type zip_file_path: str
19        """
20        self.zip_file_path = zip_file_path
21        self.zip_file = zipfile.ZipFile(zip_file_path)
22        self.tiff_file_names = [
23            name
24            for name in self.zip_file.namelist()
25            if name.endswith(".tif") and not name.startswith("_")
26        ]
27        self.tiff_file_names = sorted(self.tiff_file_names)
28        # print(self.tiff_file_names)
29
30    def __getitem__(self, index: int) -> np.ndarray:
31        if index >= len(self.tiff_file_names):
32            raise IndexError
33        else:
34            with self.zip_file.open(self.tiff_file_names[index]) as tiff_file:
35                image = skimage.io.imread(tiff_file)
36                return np.array(image)
37
38    def __len__(self) -> int:
39        return len(self.tiff_file_names)
40
41    def __iter__(self):
42        for index in range(len(self.tiff_file_names)):
43            yield self[index]
44
45    def _get_im0(self):
46        if self._im0 is None:
47            self._im0 = self[0]
48            self._shape = (
49                len(self.tiff_file_names),
50                self._im0.shape[0],
51                self._im0.shape[1],
52            )
53            self._dtype = self._im0.dtype
54        return self._im0
55
56    def get_shape(self) -> list:
57        """
58        Returns the shape of the image stack
59        :return: shape of the image stack
60        :rtype: list
61        """
62        self._get_im0()
63        return self._shape
64
65    def get_dtype(self) -> str:
66        """
67        Returns the data type of the image stack
68        :return: data type of the image stack
69        :rtype: str
70        """
71        self._get_im0()
72        return self._dtype
73
74    def get_thumb(self, save_path=None):
75        """
76        Saves a thumbnail (first frame) of the image stack
77        :param save_path: path to save the thumbnail
78        :type save_path: str
79        """
80        thumb = transform.resize(self._get_im0(), (64, 64))
81        _max = thumb.max()
82        _min = thumb.min()
83        thumb = (thumb.astype("float32") - _min / (_max - _min)) * 255
84        if save_path != None:
85            # Save the thumbnail as a JPEG file
86            skimage.io.imsave(os.path.join(save_path, "thumbnail.jpg"), thumb)
87
88    def close(self):
89        """
90        Closes the zip file
91        """
92        self.zip_file.close()